import re
import signal
from contextlib import contextmanager, redirect_stdout
from dataclasses import dataclass
from enum import Enum
from io import StringIO
from typing import Optional, Type

from ..schema import ActionReturn, ActionStatusCode
from .base_action import AsyncActionMixin, BaseAction, tool_api
from .parser import BaseParser, JsonParser


class Status(str, Enum):
    """执行状态枚举类"""
    SUCCESS = 'success'  # 执行成功
    FAILURE = 'failure'  # 执行失败


@dataclass
class ExecutionResult:
    """
    执行结果数据类
    
    属性:
        status (Status): 执行状态
        value (Optional[str]): 执行结果值
        msg (Optional[str]): 错误信息
    """
    status: Status
    value: Optional[str] = None
    msg: Optional[str] = None


@contextmanager
def _raise_timeout(timeout):
    """
    超时控制上下文管理器
    
    参数:
        timeout: 超时时间（秒）
        
    异常:
        TimeoutError: 当执行时间超过timeout时抛出
    """
    def _handler(signum, frame):
        raise TimeoutError()

    # 设置信号处理器
    signal.signal(signal.SIGALRM, _handler)
    signal.alarm(timeout)

    try:
        yield
    finally:
        signal.alarm(0)  # 清除告警


class IPythonInteractive(BaseAction):
    """
    交互式IPython shell执行器
    
    用于执行Python代码并获取执行结果
    
    参数:
        timeout (int): Python脚本执行的最大等待时间，默认30秒
        max_out_len (int): 最大输出长度，负值表示不截断，默认8192
        use_signals (bool): 是否使用信号进行超时控制，非主线程时应设为False，默认True
        description (dict): 动作描述，默认None
        parser (Type[BaseParser]): 输入输出解析器类，默认JsonParser
    """

    def __init__(
        self,
        timeout: int = 30,
        max_out_len: int = 8192,
        use_signals: bool = True,
        description: Optional[dict] = None,
        parser: Type[BaseParser] = JsonParser,
    ):
        super().__init__(description, parser)
        self.timeout = timeout
        self._executor = self.create_shell()  # 创建IPython shell实例
        # 编译ANSI转义序列的正则表达式，用于清理输出
        self._highlighting = re.compile(
            r'(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]')
        self._max_out_len = max_out_len if max_out_len >= 0 else None
        self._use_signals = use_signals

    def reset(self):
        """清除执行上下文"""
        self._executor.reset()

    @tool_api
    def run(self, command: str, timeout: Optional[int] = None) -> ActionReturn:
        """
        启动IPython交互式shell执行代码
        
        @param {string} command - 要执行的Python代码
        @param {number|undefined} timeout - 可选的超时时间（秒）
        @returns {ActionReturn} 执行结果对象
        """
        # 导入timeout装饰器并重命名为timer，用于处理代码执行超时
        from timeout_decorator import timeout as timer
        
        # 创建返回对象，记录执行命令和动作类型
        tool_return = ActionReturn(args={'text': command}, type=self.name)
        
        # 执行代码并获取结果
        # 1. 如果启用了信号处理(_use_signals为True):
        #    - 使用timer装饰器包装exec方法
        #    - timeout值使用传入的timeout或默认的self.timeout
        # 2. 如果未启用信号处理:
        #    - 直接调用exec方法执行代码
        ret = (
            timer(timeout or self.timeout)(self.exec)(command)
            if self._use_signals else self.exec(command))
            
        # 处理执行结果
        if ret.status is Status.SUCCESS:
            # 执行成功：设置结果内容和成功状态
            tool_return.result = [{'type': 'text', 'content': ret.value}]
            tool_return.state = ActionStatusCode.SUCCESS
        else:
            # 执行失败：设置错误信息和错误状态
            tool_return.errmsg = ret.msg
            tool_return.state = ActionStatusCode.API_ERROR
        
        return tool_return

    def exec(self, code: str) -> ExecutionResult:
        """
        在IPython shell中执行Python代码
        
        @param {string} code - 要执行的Python代码
        @returns {ExecutionResult} 执行结果对象
        """
        # 创建StringIO对象用于捕获输出
        output = StringIO()
        try:
            with redirect_stdout(output):
                # 执行代码并获取结果
                ret = self._executor.run_cell(self.extract_code(code))
                result = ret.result
                
                # 获取所有输出
                output_str = output.getvalue().strip()
                
                # 如果有直接返回值
                if result is not None:
                    if output_str:
                        # 如果既有输出又有返回值，将它们组合
                        combined_result = f"{output_str}\n{str(result)}"
                        return ExecutionResult(Status.SUCCESS, 
                                            combined_result[:self._max_out_len])
                    return ExecutionResult(Status.SUCCESS,
                                        str(result)[:self._max_out_len])
                                    
                # 如果只有标准输出
                if output_str:
                    outs = output_str.split('\n')
                    
                    # 检查是否有错误
                    for i, out in enumerate(outs):
                        if re.search('Error|Traceback', out, re.S):
                            if 'TimeoutError' in out:
                                return ExecutionResult(
                                    Status.FAILURE,
                                    msg='代码解释器遇到超时错误。')
                            err_idx = i
                            break
                    else:
                        return ExecutionResult(Status.SUCCESS,
                                           output_str[:self._max_out_len])
                                       
                    # 返回错误信息
                    return ExecutionResult(
                        Status.FAILURE,
                        msg=self._highlighting.sub(
                            '', '\n'.join(outs[err_idx:])[:self._max_out_len]),
                    )
                
                # 如果既没有返回值也没有输出
                return ExecutionResult(Status.SUCCESS, '')
                
        finally:
            output.close()

    @staticmethod
    def create_shell():
        """
        创建IPython shell实例，主要作用是创建一个安全的、隔离的IPython执行环境。
        
        @returns {InteractiveShell} 配置好的IPython shell实例
        """
        from IPython import InteractiveShell  # IPython的核心shell组件
        from traitlets.config import Config   # IPython的配置管理器

        # 创建配置对象
        c = Config()
        # 禁用历史记录功能，避免在磁盘上保存执行历史
        c.HistoryManager.enabled = False
        # 将历史文件设置为内存模式，进一步确保不写入磁盘
        c.HistoryManager.hist_file = ':memory:'

        # 创建并返回配置好的shell实例
        return InteractiveShell(
            # user_ns: 用户命名空间，注入超时控制函数
            user_ns={'_raise_timeout': _raise_timeout}, 
            # 应用上面的配置
            config=c)

    @staticmethod
    def extract_code(text: str) -> str:
        """
        从标记语言中提取Python代码
        
        参数:
            text (str): Markdown格式的文本
            
        返回:
            str: 提取的Python代码
        """
        import json5

        # 首先匹配三重反引号代码块
        triple_match = re.search(r'```[^\n]*\n(.+?)```', text, re.DOTALL)
        # 其次匹配单反引号代码块
        single_match = re.search(r'`([^`]*)`', text, re.DOTALL)
        
        if triple_match:
            text = triple_match.group(1)
        elif single_match:
            text = single_match.group(1)
        else:
            try:
                text = json5.loads(text)['code']
            except Exception:
                pass
        return text

    @staticmethod
    def wrap_code_with_timeout(code: str, timeout: int) -> str:
        """
        使用超时控制包装代码
        
        参数:
            code (str): 原始代码
            timeout (int): 超时时间
            
        返回:
            str: 包装后的代码
        """
        if not code.strip():
            return code
            
        code = code.strip('\n').rstrip()
        indent = len(code) - len(code.lstrip())
        # 添加超时控制
        handle = ' ' * indent + f'with _raise_timeout({timeout}):\n'
        block = '\n'.join(['    ' + line for line in code.split('\n')])
        wrapped_code = handle + block
        
        # 处理表达式语句
        last_line = code.split('\n')[-1]
        is_expression = True
        try:
            compile(last_line.lstrip(), '<stdin>', 'eval')
        except SyntaxError:
            is_expression = False
        if is_expression:
            wrapped_code += '\n' * 5 + last_line
        return wrapped_code


class AsyncIPythonInteractive(AsyncActionMixin, IPythonInteractive):
    """
    异步交互式IPython shell执行器
    
    提供异步执行Python代码的能力
    
    参数与IPythonInteractive相同
    """

    @tool_api
    async def run(self,
                  command: str,
                  timeout: Optional[int] = None) -> ActionReturn:
        """
        异步启动IPython交互式shell执行代码
        
        参数:
            command (str): Python代码片段
            timeout (Optional[int]): 执行超时时间
            
        返回:
            ActionReturn: 执行结果
        """
        tool_return = ActionReturn(args={'text': command}, type=self.name)
        ret = await self.exec(command, timeout)
        if ret.status is Status.SUCCESS:
            tool_return.result = [{'type': 'text', 'content': ret.value}]
            tool_return.state = ActionStatusCode.SUCCESS
        else:
            tool_return.errmsg = ret.msg
            tool_return.state = ActionStatusCode.API_ERROR
        return tool_return

    async def exec(self, code: str, timeout: int = None) -> ExecutionResult:
        """
        异步执行Python代码
        
        @param {string} code - 代码块
        @param {int} timeout - 执行超时时间
        @returns {ExecutionResult} 执行结果对象
        """
        output = StringIO()
        try:
            with redirect_stdout(output):
                ret = await self._executor.run_cell_async(
                    self.wrap_code_with_timeout(
                        self.extract_code(code), timeout or self.timeout))
                result = ret.result
                
                # 获取所有输出
                output_str = output.getvalue().strip()
                
                # 如果有直接返回值
                if result is not None:
                    if output_str:
                        # 如果既有输出又有返回值，将它们组合
                        combined_result = f"{output_str}\n{str(result)}"
                        return ExecutionResult(Status.SUCCESS, 
                                            combined_result[:self._max_out_len])
                    return ExecutionResult(Status.SUCCESS,
                                        str(result)[:self._max_out_len])
                                    
                # 如果只有标准输出
                if output_str:
                    outs = output_str.split('\n')
                    
                    # 检查是否有错误
                    for i, out in enumerate(outs):
                        if re.search('Error|Traceback', out, re.S):
                            if 'TimeoutError' in out:
                                return ExecutionResult(
                                    Status.FAILURE,
                                    msg='代码解释器遇到超时错误。')
                            err_idx = i
                            break
                    else:
                        return ExecutionResult(Status.SUCCESS,
                                           output_str[:self._max_out_len])
                                       
                # 返回错误信息
                return ExecutionResult(
                    Status.FAILURE,
                    msg=self._highlighting.sub(
                        '', '\n'.join(outs[err_idx:])[:self._max_out_len]),
                )
                
        finally:
            output.close()
